import librosa
import numpy as np
import soundfile as sf
from pathlib import Path
import matplotlib.pyplot as plt
from typing import Union, Tuple

## 音频处理
# 音频加载与格式统一
# Mel频谱特征提取
# 静音修剪与音频保存
# 波形/频谱可视化

def load_audio(
    path: Union[str, Path],
    target_sr: int = 44100,
    mono: bool = True
) -> Tuple[np.ndarray, int]:
    """加载音频并统一采样率"""
    y, sr = librosa.load(path, sr=target_sr, mono=mono)
    return y, sr

def extract_melspectrogram(
    y: np.ndarray,
    sr: int = 44100,
    n_fft: int = 2048,
    hop_length: int = 512,
    n_mels: int = 128,
    db_scale: bool = True
) -> np.ndarray:
    """提取Mel频谱特征"""
    mel = librosa.feature.melspectrogram(
        y=y, sr=sr,
        n_fft=n_fft,
        hop_length=hop_length,
        n_mels=n_mels
    )
    return librosa.power_to_db(mel, ref=np.max) if db_scale else mel

def trim_silence(
    y: np.ndarray,
    top_db: int = 20,
    frame_length: int = 2048,
    hop_length: int = 512
) -> np.ndarray:
    """去除静音片段"""
    y_trimmed, _ = librosa.effects.trim(
        y,
        top_db=top_db,
        frame_length=frame_length,
        hop_length=hop_length
    )
    return y_trimmed

def save_audio(
    y: np.ndarray,
    path: Union[str, Path],
    sr: int = 44100
) -> None:
    """保存音频文件"""
    sf.write(path, y, sr)

def plot_waveform(
    y: np.ndarray,
    sr: int = 44100,
    title: str = "Waveform"
) -> None:
    """绘制波形图"""
    plt.figure(figsize=(12, 4))
    librosa.display.waveshow(y, sr=sr)
    plt.title(title)
    plt.tight_layout()
    plt.show()

def plot_spectrogram(
    S: np.ndarray,
    sr: int = 44100,
    hop_length: int = 512,
    title: str = "Spectrogram"
) -> None:
    """绘制频谱图"""
    plt.figure(figsize=(12, 6))
    librosa.display.specshow(
        S,
        sr=sr,
        hop_length=hop_length,
        x_axis='time',
        y_axis='mel'
    )
    plt.colorbar(format='%+2.0f dB')
    plt.title(title)
    plt.tight_layout()
    plt.show()


## 使用示例
# from utils.audio_utils import load_audio, extract_melspectrogram
#
# y, sr = load_audio("test.wav")
# mel = extract_melspectrogram(y, sr)